#include "dataloader.h"
#include "FaceAlignment.h"
#include "mtcnn.h"
#include <algorithm>
using std::cout;
using std::endl;
using cv::Mat;

namespace filesystem = std::experimental::filesystem;

void DataLoader::save_face_image(Mat face_img, Mat body_img, const std::pair<std::string, float>& person_info, const string& file_path)
{
	filesystem::path lpath(file_path);
	if (!filesystem::exists(lpath))
		filesystem::create_directories(lpath);

	auto time_now = chrono::time_point_cast<chrono::milliseconds>(chrono::system_clock::now());
	std::time_t tt = chrono::system_clock::to_time_t(time_now);
	std::tm tm = *std::localtime(&tt);

	auto diff = chrono::duration_cast<chrono::milliseconds>(time_now.time_since_epoch()).count();
	int const msecs = diff % 1000;

	char mbstr[100] = { 0 };
	std::strftime(mbstr, sizeof(mbstr), "%Y%m%d%H%M%S", std::localtime(&tt));

	string time_str(mbstr);
	char mbstr_msec[4];
	sprintf(mbstr_msec, "%d", msecs);
	time_str += mbstr_msec;

	char score_str[10];
	sprintf(score_str, "%0.4f", person_info.second);

	imwrite(file_path + time_str + "_" + person_info.first+"_"+ score_str + "_face.jpg", face_img);
	imwrite(file_path + time_str + "_" + person_info.first + "_" + score_str  + "_body.jpg", body_img);

	/*fstream fout(file_path+"/label.txt", ios::app);
	fout << time_str + ".jpg" << " " << person_name << endl;
	fout.close();*/
}

void DataLoader::save_face_image(Mat frame, Mat detect_region, const std::vector<cv::Point2f>& landmarks, const std::pair<std::string, float>& person_info, const string& file_path)
{
	cv::Mat face_img =  fr::FaceAlignment::align_face(detect_region, landmarks, 96, 112);
	save_face_image(face_img, frame, person_info, file_path);
}

pair<std::string, vector<float>> DataLoader::process_oneline(const std::string & str)
{
	std::regex re(" ");
	std::sregex_token_iterator
		first{ str.begin(), str.end(), re, -1 },
		last;
	std::vector<string> str_line{ first, last };

	string person_id = str_line[0];
	std::vector<float> feature(str_line.size() - 1, 0.0f);
	std::transform(str_line.begin() + 1, str_line.end(), feature.begin(),
		[](const string& s)->float {return std::stof(s); });

	return make_pair(person_id, std::move(feature));
}

void DataLoader::record_oneline(std::fstream& output_file, const std::string& person_name, const vector<float>& feature)
{
	output_file << person_name;

	for (auto it = feature.begin(); it != feature.end(); ++it)
	{
		output_file << " " << std::setprecision(6) << *it;
	}

	output_file << "\n";
}

string DataLoader::get_name_from_filename(const string& image_name)
{
	auto pos = image_name.find_first_of('_');
	return image_name.substr(0, pos);
}

vector<pair<string, string>> DataLoader::load_facedb(const string& root_dir) {
	vector<pair<string, string>> facedb;
	if (!root_dir.empty())
	{
		filesystem::path apk_path(root_dir);
		filesystem::recursive_directory_iterator end;

		for (filesystem::recursive_directory_iterator i(apk_path); i != end; ++i)
		{
			if (filesystem::is_regular_file(i->path())) {
				const filesystem::path cp = (*i);
				facedb.push_back(make_pair(cp.filename().string(), cp.string()));
			}
		}
	}
	return facedb;
};

void DataLoader::load_face_db(const string & face_db_folder, const string & face_db_file, fr::FaceRecognizer * face_recognizer, const std::vector<std::string>& mtcnn_params)
{
	if (face_db_folder.empty() && face_db_file.empty())
		return;

	if(face_recognizer ==nullptr)
		std::cerr << "shoud set a face_recognizer!" << endl;

	fstream rec_file;
	std::map<std::string, vector<float>> rec_features;
	if (!face_db_file.empty())
	{
		rec_file.open(face_db_file, ios::in | ios::out | ios::app);
		if (!rec_file)
			std::cerr << "can not open database record file" << endl;

		std::string rec_line;
		while (std::getline(rec_file, rec_line))
		{
			auto&& feature = process_oneline(rec_line);

			rec_features.insert(feature);

			std::string person_name = get_name_from_filename(feature.first);
			face_recognizer->add_person(person_name, feature.second);
		}

		rec_file.clear();
	}

	if (!face_db_folder.empty())
	{
		auto facedb = load_facedb(face_db_folder);

		for (size_t i = 0; i < facedb.size(); ++i)
		{
			if (rec_features.find(facedb[i].first) != rec_features.end())
			{
				continue;
			}

			cout << facedb[i].second << endl;
			auto src = imread(facedb[i].second);

			vector<cv::Point2f> landmarks;
			std::vector<Bbox> faces;
			landmarks.reserve(5);

			faces.clear();
			landmarks.clear();
			//do face detection
			mtcnn face_detector(src.rows, src.cols, mtcnn_params);
			face_detector.detectFace(src, faces);

			if (!faces.empty())
			{
				int face_index = 0;
				if (faces.size() > 1)
				{
					//pick face closest to center
					cv::Size center = src.size() / 2;
					int min_dis = INT_MAX;
					for (size_t i = 0; i < faces.size(); i++)
					{
						cv::Point face_center = cv::Point(faces[i].y1 + faces[i].y2, faces[i].x2 + faces[i].x1);
						int dist = pow(face_center.x / 2 - center.width, 2.0) + pow(face_center.y / 2 - center.height, 2.0);

						if (dist < min_dis)
						{
							min_dis = dist;
							face_index = i;
						}
					}
				}

				//assign landmarks
				for (int i = 0; i < 5; ++i)
				{
					landmarks.push_back(cv::Point2f(faces[face_index].ppoint[i], faces[face_index].ppoint[i + 5]));
				}

				src = fr::FaceAlignment::align_face(src, landmarks, 96, 112);
			}

			auto feature = face_recognizer->extract_feature(src);
			auto name = get_name_from_filename(facedb[i].first);
			face_recognizer->add_person(name, feature);
			record_oneline(rec_file, facedb[i].first, feature);
		}

		rec_file.close();
	}
}
